#include <iostream>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgcodecs.hpp>
#include <ctime>  // 用于生成时间戳

#include "yolov6.h"

std::string getCurrentTimestamp() {
    std::time_t now = std::time(nullptr);
    char buf[64];
    std::strftime(buf, sizeof(buf), "%Y%m%d%H%M%S", std::localtime(&now));
    return std::string(buf);
}

int main(int argc, char **argv)
{
    if (argc != 3) {
        std::cout << "Usage: ./onnx-demo model_path image_path" << std::endl;
        return -1;
    }

    std::string model_path = argv[1];
    std::string image_path = argv[2];

    // 配置 YOLOv6 模型
    NetConfig YOLOV6_nets = {0.3, 0.5, model_path, "models/coco.names"};
    YOLOV6 YOLOV6_model(YOLOV6_nets);

    // 读取输入图像
    cv::Mat srcimg = cv::imread(image_path);
    if (srcimg.empty()) {
        std::cerr << "Error: Unable to read image: " << image_path << std::endl;
        return -1;
    }

    // 检测物体
    YOLOV6_model.detect(srcimg);

    // 生成输出图像路径
    std::string timestamp = getCurrentTimestamp();
    std::string output_path = image_path.substr(0, image_path.find_last_of('.')) + "_" + timestamp + ".jpg";

    // 保存带标注的图像
    if (!cv::imwrite(output_path, srcimg)) {
        std::cerr << "Error: Unable to save image: " << output_path << std::endl;
        return -1;
    }

    std::cout << "Detection result saved to: " << output_path << std::endl;

    return 0;
}
